package limiter

import (
	"sync"
	"time"
)

type Limiter struct {
	duration    int
	maxReqCount int
	mu          sync.Mutex
	timestamps  []time.Time
}

func NewLimiter(duration, maxReqCount int) *Limiter {
	return &Limiter{
		duration:    duration,
		maxReqCount: maxReqCount,
		timestamps:  make([]time.Time, 0),
	}
}

func (l *Limiter) Allow() bool {
	l.mu.Lock()
	defer l.mu.Unlock()

	now := time.Now()
	windowStart := now.Add(-time.Duration(l.duration) * time.Second)

	// 移除过期的请求
	for len(l.timestamps) > 0 && l.timestamps[0].Before(windowStart) {
		l.timestamps = l.timestamps[1:]
	}

	if len(l.timestamps) < l.maxReqCount {
		// 允许请求，并记录时间戳
		l.timestamps = append(l.timestamps, now)
		return true
	}

	// 超过限制，拒绝请求
	return false
}

func (l *Limiter) Exec(f func()) {
	for {
		if l.Allow() {
			f()
			return
		}

		l.mu.Lock()
		nextAllowedTime := l.timestamps[0].Add(time.Duration(l.duration) * time.Second)
		l.mu.Unlock()

		// 等待直到下一个允许的时间点
		<-time.After(time.Until(nextAllowedTime))
	}
}

func (l *Limiter) Wait() {
	for {
		if l.Allow() {
			return
		}

		l.mu.Lock()
		nextAllowedTime := l.timestamps[0].Add(time.Duration(l.duration) * time.Second)
		l.mu.Unlock()

		// 等待直到下一个允许的时间点
		<-time.After(time.Until(nextAllowedTime))
	}
}
